from typing import Dict, List

from lagent.memory import Memory
from lagent.prompts import StrParser


class DefaultAggregator:
    """
    默认消息聚合器
    
    负责将历史消息和系统指令聚合成标准格式，用于LLM的输入
    """

    def aggregate(self,
                 messages: Memory,
                 name: str,
                 parser: StrParser = None,
                 system_instruction: str = None) -> List[Dict[str, str]]:
        """
        聚合消息和系统指令
        
        参数:
            messages (Memory): 历史消息记忆
            name (str): 代理名称
            parser (StrParser): 消息解析器，可选
            system_instruction (str): 系统指令，可选
            
        返回:
            List[Dict[str, str]]: 聚合后的消息列表，每个消息包含role和content
        """
        _message = []  # 存储聚合后的消息
        messages = messages.get_memory()  # 获取历史消息
        
        # 添加系统指令（如果有）
        if system_instruction:
            _message.extend(
                self.aggregate_system_intruction(system_instruction))
                
        # 处理历史消息
        for message in messages:
            if message.sender == name:
                # 处理代理发送的消息
                _message.append(
                    dict(role='assistant', content=str(message.content)))
            else:
                # 处理用户发送的消息
                user_message = message.content
                # 如果上一条也是用户消息，则合并内容
                if len(_message) > 0 and _message[-1]['role'] == 'user':
                    _message[-1]['content'] += user_message
                else:
                    _message.append(dict(role='user', content=user_message))
        return _message

    @staticmethod
    def aggregate_system_intruction(system_intruction) -> List[dict]:
        """
        聚合系统指令
        
        将系统指令转换为标准格式
        
        参数:
            system_intruction: 系统指令，可以是字符串、字典或列表
            
        返回:
            List[dict]: 格式化后的系统指令列表
            
        异常:
            TypeError: 当消息格式不支持时抛出
            KeyError: 当缺少必需的键时抛出
        """
        # 处理字符串格式的系统指令
        if isinstance(system_intruction, str):
            system_intruction = dict(role='system', content=system_intruction)
            
        # 处理字典格式的系统指令
        if isinstance(system_intruction, dict):
            system_intruction = [system_intruction]
            
        # 验证系统指令格式
        if isinstance(system_intruction, list):
            for msg in system_intruction:
                # 检查消息类型
                if not isinstance(msg, dict):
                    raise TypeError(f'Unsupported message type: {type(msg)}')
                # 检查必需的键
                if not ('role' in msg and 'content' in msg):
                    raise KeyError(
                        f"Missing required key 'role' or 'content': {msg}")
                        
        return system_intruction
